import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
from sklearn.feature_extraction.text import TfidfVectorizer
import lightgbm as lgb
import warnings
import time
from tqdm import tqdm

warnings.filterwarnings('ignore')


# 1. 数据加载和预处理
def load_data(train_path, test_path):
    print("Loading data...")
    train_df = pd.read_csv(train_path, sep='\t')
    test_df = pd.read_csv(test_path, sep='\t')

    # 将文本转换为字符串类型
    train_df['text'] = train_df['text'].astype(str)
    test_df['text'] = test_df['text'].astype(str)

    # 将数字序列转换为空格分隔的字符串
    train_df['text'] = train_df['text'].apply(lambda x: ' '.join(x.split()))
    test_df['text'] = test_df['text'].apply(lambda x: ' '.join(x.split()))

    print(f"Train shape: {train_df.shape}")
    print(f"Test shape: {test_df.shape}")

    return train_df, test_df


# 2. 特征工程
def feature_engineering(train_df, test_df):
    print("Feature engineering...")
    tfidf = TfidfVectorizer(
        ngram_range=(1, 2),
        max_features=20000,
        min_df=3
    )

    print("Fitting TF-IDF...")
    train_features = tfidf.fit_transform(train_df['text'])
    print("Transforming test set...")
    test_features = tfidf.transform(test_df['text'])

    print(f"Train features shape: {train_features.shape}")
    print(f"Test features shape: {test_features.shape}")

    return train_features, test_features


# 3. 模型训练
def train_model(X_train, X_val, y_train, y_val):
    print("Training model...")
    params = {
        'objective': 'multiclass',
        'num_class': 14,
        'metric': 'multi_logloss',
        'learning_rate': 0.3,  # 进一步增加学习率
        'num_leaves': 31,
        'max_depth': 5,  # 进一步减小树的深度
        'feature_fraction': 0.8,
        'bagging_fraction': 0.8,
        'bagging_freq': 5,
        'verbose': -1
    }

    # 创建数据集
    train_data = lgb.Dataset(X_train, label=y_train)
    val_data = lgb.Dataset(X_val, label=y_val)

    # 训练模型
    callbacks = [
        lgb.early_stopping(stopping_rounds=5),
        lgb.log_evaluation(period=1)
    ]

    print("Starting training...")
    model = lgb.train(
        params,
        train_data,
        num_boost_round=50,
        valid_sets=[val_data],
        callbacks=callbacks
    )

    return model


def save_submission(predictions, filename='submission.csv'):
    """保存预测结果到文件"""
    submission = pd.DataFrame({
        'label': predictions
    })
    submission.to_csv(filename, index=False)
    print(f"Predictions saved to {filename}")


def main():
    # 设置文件路径（当前目录）
    train_path = 'train_set.csv'
    test_path = 'test_a.csv'

    try:
        # 1. 加载数据
        print("Step 1: Loading data...")
        start_time = time.time()
        train_df, test_df = load_data(train_path, test_path)
        print(f"Data loading completed in {time.time() - start_time:.2f} seconds\n")

        # 2. 特征工程
        print("Step 2: Feature engineering...")
        start_time = time.time()
        train_features, test_features = feature_engineering(train_df, test_df)
        print(f"Feature engineering completed in {time.time() - start_time:.2f} seconds\n")

        # 3. 分割训练集和验证集
        print("Step 3: Splitting data...")
        start_time = time.time()
        X_train, X_val, y_train, y_val = train_test_split(
            train_features, train_df['label'],
            test_size=0.2,
            random_state=42,
            stratify=train_df['label']
        )
        print(f"Training set size: {X_train.shape}")
        print(f"Validation set size: {X_val.shape}")
        print(f"Data splitting completed in {time.time() - start_time:.2f} seconds\n")

        # 4. 训练模型
        print("Step 4: Training model...")
        start_time = time.time()
        model = train_model(X_train, X_val, y_train, y_val)
        print(f"Model training completed in {time.time() - start_time:.2f} seconds\n")

        # 5. 验证集预测
        print("Step 5: Evaluating model...")
        start_time = time.time()
        val_pred = model.predict(X_val)
        val_pred_labels = np.argmax(val_pred, axis=1)
        val_f1 = f1_score(y_val, val_pred_labels, average='macro')
        print(f"Validation Macro F1 Score: {val_f1:.4f}")
        print(f"Evaluation completed in {time.time() - start_time:.2f} seconds\n")

        # 6. 测试集预测
        print("Step 6: Predicting test set...")
        start_time = time.time()
        test_pred = model.predict(test_features)
        test_pred_labels = np.argmax(test_pred, axis=1)
        print(f"Prediction completed in {time.time() - start_time:.2f} seconds\n")

        # 7. 保存结果
        print("Step 7: Saving predictions...")
        save_submission(test_pred_labels)

        print("\nAll steps completed successfully!")

    except Exception as e:
        print(f"An error occurred: {str(e)}")
        raise


if __name__ == "__main__":
    main()
